package top.yudoge.dogeeditorserverj.llm;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.agent.tool.ToolSpecifications;
import dev.langchain4j.data.message.*;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import lombok.Data;
import org.springframework.stereotype.Component;
import top.yudoge.dogeeditorserverj.llm.tools.DataBaseTools;
import top.yudoge.dogeeditorserverj.transients.query.RewriteQuery;

import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;

@Component
public class SQLRewriteAgent implements RewriteAgent {

    @Override
    public void doRewrite(RewriteQuery query, LLMStreamConsumer consumer) {
        SQLLangParams sqlLangParams = verifyAndPreHandle(query);

        // 1. Reasoning the task
        ChatLanguageModel reasonModel = LLMFactory.getReasonModel(query.getModelName());
        ChatRequest request = ChatRequest.builder()
                .messages(messages(query, sqlLangParams))
                .parameters(
                        ChatRequestParameters.builder()
                                .toolSpecifications(getTools())
                                .build()
                ).build();

        do {
            ChatResponse resp = reasonModel.chat(request);
            AiMessage aiMessage = resp.aiMessage();
            System.out.println(aiMessage.text());
            if (aiMessage.hasToolExecutionRequests()) {
                List<ToolExecutionRequest> toolExecutionRequests = aiMessage.toolExecutionRequests();
                for (ToolExecutionRequest toolExecutionRequest : toolExecutionRequests) {
                    System.out.println(toolExecutionRequest.name());
                    System.out.println(toolExecutionRequest.arguments());
                    System.out.println(toolExecutionRequest.id());
                }

            } else {
                System.out.println(aiMessage.text());
            }

        } while (true);

    }


    private List<ToolSpecification> getTools() {
        return ToolSpecifications.toolSpecificationsFrom(DataBaseTools.class);
    }

    private List<ChatMessage> messages(RewriteQuery query, SQLLangParams sqlLangParams) {
        List<ChatMessage> messages = new ArrayList<>();
        // build system message
        SystemMessage systemMessage = new SystemMessage("" +
                "You are an excellent SQL engineer. Your task is helping users rewrite their SQL script.\n" +
                "1. You should think step by step. I'll give you some tools that you may use in some step.\n" +
                "2. You may receive a part of user code. Your task is rewrite this.\n" +
                "3. If there's no user code. You should write it from scratch.\n" +
                "4. I'll give you some information such as the datasource, database, schema that the user currently in. You maybe use it when calling the tools."
        );

        messages.add(systemMessage);

        StringBuffer sb = new StringBuffer();
        sb.append("The dataSourceName, databaseName, schemaName and the table user selected now is => ");
        sb.append(JSONUtil.toJsonStr(sqlLangParams));
        sb.append("\n");
        UserMessage dbSelectedMessage = new UserMessage(sb.toString());

        messages.add(dbSelectedMessage);

        if (!StrUtil.isBlank(query.getEditorCode())) {
            UserMessage userCodeMessage = new UserMessage("User code is : \n" + query.getEditorCode());
            messages.add(userCodeMessage);
        }

        UserMessage userTaskMessage = new UserMessage("User rewrite task is : \n" + query.getUserTask());
        messages.add(userTaskMessage);

        return messages;
    }




    private SQLLangParams verifyAndPreHandle(RewriteQuery query) {
        query.setModelName("qwenStreamingChatModel");
        SQLLangParams sqlLangParams = BeanUtil.toBean(query.getLangParams(), SQLLangParams.class);
        if (Objects.isNull(sqlLangParams)) {
            throw new RewriteQueryInvalidateException("You must specified dsName, dbName, schemaName in langParams");
        }
        if (StrUtil.isBlank(sqlLangParams.dataSourceName)) {
            throw new RewriteQueryInvalidateException("You must specified dsName, dbName, schemaName in langParams");
        }
        if (StrUtil.isBlank(sqlLangParams.databaseName)) {
            throw new RewriteQueryInvalidateException("You must specified dsName, dbName, schemaName in langParams");
        }
        if (StrUtil.isBlank(sqlLangParams.schemaName)) {
            throw new RewriteQueryInvalidateException("You must specified dsName, dbName, schemaName in langParams");
        }
        if (Objects.isNull(sqlLangParams.tableSelected)) {
            sqlLangParams.tableSelected = Collections.emptyList();
        }

        return sqlLangParams;
    }

    @Data
    private static class SQLLangParams {
        private String dataSourceName;
        private String databaseName;
        private String schemaName;
        private List<String> tableSelected;
    }

    private String nullTo(String str, String defaultStr) {
        if (str == null) {
            return defaultStr;
        }
        return str;
    }
}
